import sys

from PyQt5.QtWidgets import QApplication
from PyQt5.QtWidgets import QWidget
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtGui import QImage
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import QTimer
import cv2
import numpy as np

from Ui_ui import *

class MainWindow(QtWidgets.QMainWindow):
    # class constructor
    def __init__(self):
        # call QWidget constructor
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        
        # create a timer
        self.timer = QTimer()
        self.timer.timeout.connect(self.viewCam)
        self.ui.action_open.triggered.connect(self.controlTimer)
        self.ui.hSlider_type.valueChanged.connect(self.threshold_type_change)
        self.ui.action_save.triggered.connect(self.saveslot)
        self.ui.action_restore.triggered.connect(self.restore_all)
        
    # view camera
    def viewCam(self):
        ret, self.im = self.cap.read()
        image = cv2.cvtColor(self.im, cv2.COLOR_BGR2RGB)
        height, width, channel = image.shape
        step = channel * width
        qImg = QImage(image.data, width, height, step, QImage.Format_RGB888)
        self.ui.label.setPixmap(QPixmap.fromImage(qImg))
        self.convert()

    def convert(self):
        if not self.ui.action_gray.isChecked() and not self.ui.action_binary.isChecked():
            self.output = cv2.cvtColor(self.im, cv2.COLOR_BGR2RGB)
            height, width, channel= self.output.shape
            step = width * channel
            isRGB = True

        if self.ui.action_gray.isChecked() or self.ui.action_binary.isChecked():
            image2 = cv2.cvtColor(self.im, cv2.COLOR_RGB2GRAY)
            height, width = image2.shape
            step = width
            isRGB = False

            if self.ui.action_binary.isChecked():
                ret, self.output =  cv2.threshold(image2, self.ui.hSlider_threshold.value(), 255, self.ui.hSlider_type.value())
            else:
                self.output = image2

        if self.ui.action_horizontal_flip.isChecked():
            self.output = cv2.flip(self.output, 1)

        if self.ui.action_verital_flip.isChecked():
             self.output = cv2.flip(self.output, 0)           

        if self.ui.action_affine.isChecked():
            srcTri = np.array( [[0, 0], [self.output.shape[1] - 1, 0], [0, self.output.shape[0] - 1]] ).astype(np.float32)
            dstTri = np.array( [[0, self.output.shape[1]*0.33], [self.output.shape[1]*0.85, self.output.shape[0]*0.25], [self.output.shape[1]*0.15, self.output.shape[0]*0.7]] ).astype(np.float32)
            warp_mat = cv2.getAffineTransform(srcTri, dstTri)
            warp_dst = cv2.warpAffine(self.output, warp_mat, (self.output.shape[1], self.output.shape[0]))
            # Rotating the image after Warp
            center = (warp_dst.shape[1]//2, warp_dst.shape[0]//2)
            angle = self.ui.hSlider_angle.value()
            scale = 0.6
            rot_mat = cv2.getRotationMatrix2D(center, angle, scale)
            self.output = cv2.warpAffine(warp_dst, rot_mat, (warp_dst.shape[1], warp_dst.shape[0]))

        if self.ui.action_erosion.isChecked():
            erosion_size = self.ui.hSlider_erosion_size.value()
            element = cv2.getStructuringElement(cv2.MORPH_RECT, (2*erosion_size + 1, 2*erosion_size+1), (erosion_size, erosion_size))
            self.output = cv2.erode(self.output, element)

        if self.ui.action_dilatation.isChecked():
            dilatation_size = self.ui.hSlider_dilatation_size.value()
            element = cv2.getStructuringElement(cv2.MORPH_RECT, (2*dilatation_size + 1, 2*dilatation_size+1), (dilatation_size, dilatation_size))
            self.output = cv2.dilate(self.output, element)

        if self.ui.action_sobel.isChecked():
            temp = cv2.GaussianBlur(self.output, (3, 3), 0)
            grad_x = cv2.Sobel(self.output, cv2.CV_16S, 1, 0, ksize=3, scale=1, delta=0, borderType=cv2.BORDER_DEFAULT)
            grad_y = cv2.Sobel(self.output, cv2.CV_16S, 0, 1, ksize=3, scale=1, delta=0, borderType=cv2.BORDER_DEFAULT)
            abs_grad_x = cv2.convertScaleAbs(grad_x)
            abs_grad_y = cv2.convertScaleAbs(grad_y)
            self.output = cv2.addWeighted(abs_grad_x, 0.5, abs_grad_y, 0.5, 0)

        if self.ui.action_laplacian.isChecked():
            temp = cv2.GaussianBlur(self.output, (3, 3), 0)
            dst = cv2.Laplacian(temp, cv2.CV_16S, ksize=3)
            self.output = cv2.convertScaleAbs(dst)

        if self.ui.action_canny.isChecked():
            ratio = 3
            kernel_size = 3
            low_threshold = self.ui.hlSlider_canny_threshold.value()
            temp = cv2.blur(self.output, (3,3))
            detected_edges = cv2.Canny(temp, low_threshold, low_threshold*ratio, kernel_size)
            mask = detected_edges != 0
            self.output = self.im * (mask[:,:,None].astype(self.im.dtype))

        if self.ui.action_morph_open.isChecked():
            erosion_size = self.ui.hSlider_erosion_size.value()            
            dilatation_size = self.ui.hSlider_dilatation_size.value()
            element1 = cv2.getStructuringElement(cv2.MORPH_RECT, (2*erosion_size + 1, 2*erosion_size+1), (erosion_size, erosion_size))
            element2 = cv2.getStructuringElement(cv2.MORPH_RECT, (2*dilatation_size + 1, 2*dilatation_size+1), (dilatation_size, dilatation_size))
            self.output = cv2.erode(self.output, element1)
            self.output = cv2.dilate(self.output, element2)

        if self.ui.action_morph_close.isChecked():
            erosion_size = self.ui.hSlider_erosion_size.value()            
            dilatation_size = self.ui.hSlider_dilatation_size.value()
            element1 = cv2.getStructuringElement(cv2.MORPH_RECT, (2*erosion_size + 1, 2*erosion_size+1), (erosion_size, erosion_size))
            element2 = cv2.getStructuringElement(cv2.MORPH_RECT, (2*dilatation_size + 1, 2*dilatation_size+1), (dilatation_size, dilatation_size))
            self.output = cv2.dilate(self.output, element2)
            self.output = cv2.erode(self.output, element1)


        if isRGB:
            qImg2 = QImage(self.output.data, width, height, step, QImage.Format_RGB888)
        else:
            qImg2 = QImage(self.output.data, width, height, step, QImage.Format_Grayscale8)
        self.ui.label_2.setPixmap(QPixmap.fromImage(qImg2))

    # start/stop timer
    def controlTimer(self):
        # if timer is stopped
        if not self.timer.isActive():
            # create video capture
            self.cap = cv2.VideoCapture(0)
            # start timer
            self.timer.start(50)
        # if timer is started
        else:
            # stop timer
            self.timer.stop()
            # release video capture
            self.cap.release()

    def threshold_type_change(self):
        threshold_type = self.ui.hSlider_type.value()
        if threshold_type == 0:
            self.ui.label_3.setText("二值化方法：cv.THRESH_BINARY")
        if threshold_type == 1:
            self.ui.label_3.setText("二值化方法：cv.THRESH_BINARY_INV")
        if threshold_type == 2:
            self.ui.label_3.setText("二值化方法：cv.THRESH_TRUNC")
        if threshold_type == 3:
            self.ui.label_3.setText("二值化方法：cv.THRESH_TOZERO")
        if threshold_type == 4:
            self.ui.label_3.setText("二值化方法：cv.THRESH_TOZERO_INV")
    
    def saveslot(self):
        fileName, tmp = QFileDialog.getSaveFileName(
            self, 'Save Image', './__data', '*.png *.jpg *.bmp', '*.png')
        if fileName is '':
            return
        if self.output.size == 1:
            return
        cv2.imwrite(fileName, self.output)

    def restore_all(self):
        self.ui.action_gray.setChecked(False)
        self.ui.action_binary.setChecked(False)
        self.ui.action_erosion.setChecked(False)
        self.ui.action_dilatation.setChecked(False)
        self.ui.action_sobel.setChecked(False)
        self.ui.action_laplacian.setChecked(False)
        self.ui.action_canny.setChecked(False)
        self.ui.action_affine.setChecked(False)

if __name__ == '__main__':
    app = QApplication(sys.argv)

    # create and show mainWindow
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())